
from datetime import date, datetime
from decimal import ROUND_HALF_UP, Decimal, getcontext
from typing import Type
from cachetools import TTLCache, cached
from fastapi.concurrency import run_in_threadpool
import pandas as pd
from module_stock.entity.vo.kline_boll_vo import KlineBollModel
from module_stock.entity.vo.boll_divergence_vo import DivergenceModel
from utils.log_util import logger
import numpy as np
from redis import Redis
import talib  # 用于技术指标计算
from module_stock.dao.daily_dao import DailyDao
from sqlalchemy.ext.asyncio import AsyncSession
from module_stock.dao.daily_boll_dao import DailyBollDao
from module_stock.dao.kline_boll_dao import KlineBollDao
from module_stock.dao.monthly_dao import MonthlyDao
from module_stock.dao.monthly_boll_dao import MonthlyBollDao
from module_stock.dao.weekly_dao import WeeklyDao
from module_stock.dao.weekly_boll_dao import WeeklyBollDao
from module_stock.entity.do.daily_boll_do import StockDailyBoll
from module_stock.entity.do.monthly_boll_do import StockMonthlyBoll
from module_stock.entity.do.weekly_boll_do import StockWeeklyBoll
from utils.stock_util import CommonUtil

class KlineBollService():

    boll_timeperiod = 20
    boll_nbdevup = 2
    boll_nbdevdn = 2
    boll_matype = 0
    boll_redis_expire = 86400

    daily_bolls_cache : pd.DataFrame = pd.DataFrame()
    weekly_bolls_cache : pd.DataFrame = pd.DataFrame()
    monthly_bolls_cache : pd.DataFrame = pd.DataFrame()

    ttl_cache = TTLCache(maxsize=100000, ttl=boll_redis_expire)

    x_types = [
        {'data_type': 'close','lower_type': 'close', 'upper_type': 'close'},
        {'data_type': 'close','lower_type': 'low', 'upper_type': 'high'}
    ]

    @classmethod  
    async def init_boll_cache(cls, query_db: AsyncSession, kline_period:str=None):  
        if not kline_period:
            cls.daily_bolls_cache = await KlineBollService.get_bolls_pd(query_db,None,None,'daily')
            cls.weekly_bolls_cache = await KlineBollService.get_bolls_pd(query_db,None,None,'weekly')
            cls.monthly_bolls_cache = await KlineBollService.get_bolls_pd(query_db,None,None,'monthly') 
        else :
            match(kline_period):
                case 'daily':
                    cls.daily_bolls_cache = await KlineBollService.get_bolls_pd(query_db,None,None,'daily')
                case 'weekly':
                    cls.weekly_bolls_cache = await KlineBollService.get_bolls_pd(query_db,None,None,'weekly')
                case 'monthly':
                    cls.monthly_bolls_cache = await KlineBollService.get_bolls_pd(query_db,None,None,'monthly')

    @classmethod
    @cached(ttl_cache)
    def get_bolls_from_cache(cls, symbol:str, start_date: date, kline_period: str = None, 
                             data_type:str= None,lower_type:str= None, upper_type:str= None):
        df = pd.DataFrame()
        match(kline_period):
            case 'daily':
                df = cls.daily_bolls_cache 
            case 'weekly':
                df = cls.weekly_bolls_cache 
            case 'monthly':
                df = cls.monthly_bolls_cache 
        if symbol:
            df = df.loc[df['symbol'] == symbol]
        if start_date:
            df = df.loc[df['trade_date'] >= start_date]
        if data_type:
            df = df.loc[df['data_type'] == data_type]
        if lower_type:
            df = df.loc[df['lower_type'] == lower_type]
        if upper_type:
            df = df.loc[df['upper_type'] == upper_type]
        return df.to_dict('records')

    @classmethod  
    def clear_boll_cache(cls, kline_period:str=None):  
        if not kline_period:
            cls.daily_bolls_cache= pd.DataFrame()
            cls.weekly_bolls_cache = pd.DataFrame()
            cls.monthly_bolls_cache = pd.DataFrame()
        else :
            match(kline_period):
                case 'daily':
                    cls.daily_bolls_cache= pd.DataFrame()
                case 'weekly':
                    cls.weekly_bolls_cache = pd.DataFrame()
                case 'monthly':
                    cls.monthly_bolls_cache = pd.DataFrame()


    @classmethod
    def get_kline_caches(cls, kline_period:str) :
        match(kline_period):
            case 'daily':
                return cls.daily_bolls_cache
            case 'weekly':
                return cls.weekly_bolls_cache
            case 'monthly':
                return cls.monthly_bolls_cache

    @classmethod
    def remove_boll_cache(cls, symbol: str, kline_period = 'daily'):
        df = cls.get_kline_caches(kline_period)
        if not df.empty:
            df.drop(df[df['symbol'] == symbol].index, inplace=True)

    @classmethod
    def add_boll_cache(cls, add_data: pd.DataFrame, symbol: str, kline_period = 'daily'):
        match(kline_period):
            case 'daily':
                cls.daily_bolls_cache = pd.concat([cls.daily_bolls_cache, add_data], axis=0)
            case 'weekly':
                cls.weekly_bolls_cache = pd.concat([cls.weekly_bolls_cache, add_data], axis=0)
            case 'monthly':
                cls.monthly_bolls_cache = pd.concat([cls.monthly_bolls_cache, add_data], axis=0)

    @classmethod
    async def get_kline_data_pd(cls, query_db: AsyncSession, symbol: str, kline_period:str) :
        match(kline_period):
            case 'daily':
                return  (await DailyDao.get_all_daily_pd(query_db, symbol))
            case 'weekly':
                return  (await WeeklyDao.get_all_weekly_pd(query_db, symbol))
            case 'monthly':
                return  (await MonthlyDao.get_all_monthly_pd(query_db, symbol))
            
    @classmethod
    def get_kline_boll_dao(cls, kline_period:str)  -> Type[KlineBollDao]:
        match(kline_period):
            case 'daily':
                return  DailyBollDao
            case 'weekly':
                return  WeeklyBollDao
            case 'monthly':
                return  MonthlyBollDao
    
    @classmethod
    def get_kline_boll_do_model(cls, kline_period:str) :
        match(kline_period):
            case 'daily':
                return  StockDailyBoll
            case 'weekly':
                return  StockWeeklyBoll
            case 'monthly':
                return  StockMonthlyBoll
            
    @classmethod
    def del_cur_period_boll(cls, kline_period:str, bolls:list) :
        if not bolls:
            return
        match(kline_period):
            case 'daily':
                return 
            case 'weekly':
                if CommonUtil.is_date_in_current_week(bolls[-1]['trade_date']) :
                    bolls.pop()
            case 'monthly':
                if CommonUtil.is_date_in_current_month(bolls[-1]['trade_date']) :
                    bolls.pop()

    @classmethod
    async def init_boll(cls, query_db: AsyncSession, symbol: str, data_type:str, lower_type:str,
                        upper_type:str, kline_period = 'daily', use_cache : bool = False):
        if use_cache:
            cls.remove_boll_cache(symbol, kline_period)

        # 只存储最近的19条数据
        selected_columns = ['trade_date', 'symbol', 'open','close',
                            'low', 'high', 'upper', 'middle', 'lower'] 
        
        transform_colums = ['open','close', 'low', 'high', 'upper', 'middle', 'lower']

        """初始化计算BOLL指标"""
        stock_data = await cls.get_kline_data_pd(query_db, symbol, kline_period)
        if stock_data.empty or len(stock_data) < cls.boll_timeperiod:
            return

        last_19_rows, dict_records = await run_in_threadpool(cls.calc_boll, data_type, lower_type, upper_type, selected_columns, transform_colums, stock_data)
        if use_cache:
            cls.add_boll_cache(last_19_rows, symbol, kline_period)
        await cls.get_kline_boll_dao(kline_period).batch_add_boll(query_db, dict_records)
        await query_db.commit()

    @classmethod
    def calc_boll(cls, data_type, lower_type, upper_type, selected_columns, transform_colums, stock_data):
        stock_data['upper'], stock_data['middle'], stock_data['lower'] = talib.BBANDS(
            stock_data[data_type].astype(np.float64), timeperiod=cls.boll_timeperiod, 
            nbdevup=cls.boll_nbdevup, nbdevdn=cls.boll_nbdevdn, matype=cls.boll_matype)
        
        stock_data['data_type'] = data_type
        selected_columns.append('data_type')

        if lower_type:
            stock_data['bottom_divergence'] = np.where(
                stock_data[lower_type] <= stock_data['lower'],  # 低于低轨
                1,  # 信号触发
                0   # 无信号
            )
            stock_data['lower_type'] = lower_type
            selected_columns.append('lower_type')
            selected_columns.append('bottom_divergence')

        if upper_type:
            stock_data['top_divergence'] = np.where(
                stock_data[upper_type] >= stock_data['upper'],  # 高于高轨
                1,  # 信号触发
                0   # 无信号
            )
            stock_data['upper_type'] = upper_type
            selected_columns.append('upper_type')
            selected_columns.append('top_divergence')
        
        last_41_rows = stock_data[selected_columns].tail(41).fillna(0)
        last_41_rows[transform_colums] = last_41_rows[transform_colums].map(
            lambda x: Decimal(str(x)).quantize(Decimal('0.01'), rounding=ROUND_HALF_UP)
        )
        dict_records = last_41_rows.to_dict('records')
        return last_41_rows,dict_records


    @classmethod
    async def detect_realtime_divergence(cls, db: AsyncSession, symbol: str, 
                                        new_data: dict,
                                        redis, kline_period, data_type:str,lower_type:str, upper_type:str):
        """
        检测实时数据是否发生背离
        :param new_data: (date, open, close, low, high)
        """
        try :
            result = await run_in_threadpool(cls.calc_realtime_divergence, symbol, new_data, kline_period, data_type, lower_type, upper_type)
            if not result :
                return
            trade_date, record, bottom_divergence, top_divergence = result
            
            if bottom_divergence : 
                await redis.set(f'boll:rt:{trade_date}:{kline_period}:bottom_divergence:{symbol}.{data_type}.{lower_type}', record.model_dump_json(), ex=cls.boll_redis_expire) 

            if top_divergence : 
                await redis.set(f'boll:rt:{trade_date}:{kline_period}:top_divergence:{symbol}.{data_type}.{upper_type}', record.model_dump_json(), ex=cls.boll_redis_expire) 
        except Exception as e :
            logger.exception(e)

    @classmethod
    def calc_realtime_divergence(cls, symbol, new_data, kline_period, data_type, lower_type, upper_type):
        bolls = cls.get_bolls_from_cache(symbol,None, kline_period, data_type, lower_type, upper_type)
        cls.del_cur_period_boll(kline_period, bolls)      
        if not bolls or len(bolls) < cls.boll_timeperiod - 1:            
                # await cls.init_boll(db, symbol, kline_period, True)
                # bolls  =   cls.get_bolls_from_cache(symbol,None, kline_period)
                # if not bolls:   
            logger.error(f"股票:{symbol}周期:{kline_period}的boll数据未初始化")
            return 
            
        # 准备增量数据
        trade_date = new_data['trade_date']
        open = new_data['open']
        close = new_data['close']
        low = new_data['low']
        high = new_data['high']

        prices  = [row[data_type] for row in bolls]
        prices.append(new_data[data_type])
        
        uppers, middles, lowers = talib.BBANDS(
            np.array(prices, dtype=np.float64),
            timeperiod=cls.boll_timeperiod,
            nbdevup=cls.boll_nbdevup,
            nbdevdn=cls.boll_nbdevdn,
            matype=cls.boll_matype
        )
        upper = Decimal(str(uppers[-1])).quantize(Decimal('0.01'), rounding=ROUND_HALF_UP)
        middle = Decimal(str(middles[-1])).quantize(Decimal('0.01'), rounding=ROUND_HALF_UP)
        lower = Decimal(str(lowers[-1])).quantize(Decimal('0.01'), rounding=ROUND_HALF_UP)
        # 准备记录
        record = KlineBollModel()
        record.symbol = symbol
        record.trade_date = trade_date
        record.open = open
        record.close = close
        record.low = low
        record.high = high
        record.upper = upper
        record.middle = middle 
        record.lower =  lower
        record.kline_period = kline_period
        record.data_type = data_type
        record.lower_type = lower_type
        record.upper_type = upper_type
            
        bottom_divergence = False
        if lower_type:
                # 检测底背离
            bottom_divergence = new_data[lower_type] <= lower
            record.bottom_divergence = bottom_divergence  

        top_divergence = False
        if upper_type:
                # 检测顶背离
            top_divergence = new_data[upper_type] >= upper
            record.top_divergence = top_divergence
        return trade_date,record,bottom_divergence,top_divergence
    
    @classmethod
    async def clear_bolls(cls, db: AsyncSession, kline_period: str):
        await cls.get_kline_boll_dao(kline_period).clear_boll_dao(db)
        await db.commit()     

    @classmethod
    async def get_bolls(cls, db: AsyncSession, symbol:str, start_date: date, kline_period: str):
        result = await cls.get_kline_boll_dao(kline_period).get_bolls(db, symbol, start_date)
        return result
    
    @classmethod
    async def get_bolls_pd(cls, db: AsyncSession, symbol:str, start_date: date, kline_period: str):
        result = await cls.get_kline_boll_dao(kline_period).get_bolls_pd(db, symbol, start_date)
        return result
    
    @classmethod
    async def get_realtime_divergences(cls, redis:Redis):
        today = datetime.now().date()
        divergenceModel = DivergenceModel()
        today_keys = await redis.keys(f'boll:rt:{today}*')
        if today_keys :
            today_values =  await redis.mget(today_keys)
            modes = [KlineBollModel.model_validate_json(today_value) for today_value in today_values]

            divergenceModel.daily_bottoms += [model for model in modes if model.kline_period == 'daily' and model.bottom_divergence == True]
            divergenceModel.weekly_bottoms += [model for model in modes if model.kline_period == 'weekly' and model.bottom_divergence == True]
            divergenceModel.monthly_bottoms += [model for model in modes if model.kline_period == 'monthly' and model.bottom_divergence == True]
            divergenceModel.daily_tops +=  [model for model in modes if model.kline_period == 'daily' and model.top_divergence == True]
            divergenceModel.weekly_tops += [model for model in modes if model.kline_period == 'weekly' and model.top_divergence == True]
            divergenceModel.monthly_tops += [model for model in modes if model.kline_period == 'monthly' and model.top_divergence == True]
            
        return divergenceModel